// (C) 2022 Intel Corporation. All rights reserved.
// Your use of Intel Corporation's design tools, logic functions and other 
// software and tools, and its AMPP partner logic functions, and any output 
// files from any of the foregoing (including device programming or simulation 
// files), and any associated documentation or information are expressly subject 
// to the terms and conditions of the Intel Program License Subscription 
// Agreement, Intel FPGA IP License Agreement, or other applicable 
// license agreement, including, without limitation, that your use is for the 
// sole purpose of programming logic devices manufactured by Intel and sold by 
// Intel or its authorized distributors.  Please refer to the applicable 
// agreement for further details.

module A_cse_ocs_dpa_aes10c_unit #(
   parameter OCS_AES_DPA_USE_CTECH = 0,
   //Added parameter for AES', for now only used for DVB-CSA3
   parameter AES_P     = 1'b0
)
(
  input  logic clk,                 // system clock
  input  logic reset_n,             // system reset

  input  logic [255:0] aes_key_in,           // Key input bus from system
  input  logic [127:0] aes_iv_in,            // IV input bus from system
  input  logic [127:0] aes_data_in,          // Data input bus from system
  input  logic [255:0] aes_mask_seed,	     // Mask seed from system bus
  input  logic 	       aes_load_mask,	     // Load seed into the PRNG
  input  logic	       mask_logic_en,	     // Enable masking for DPA
  output logic         data_valid,            // Start AES operation
  input  logic         start,
  input  logic         aes_op,               // 0 - Encrypt, 1 - Decrypt, 2 - Expand, 3 - Bypass
  input  logic         aes_load_iv,          // Load the IV into the working buffer
  input  logic         aes_flush,            // Flush all internal flops that could store data related to previous keys/data
  input  logic [2:0]   aes_internal_mode,    // AES Mode: ECB = 3'b000, CBC = 3'b001, CTR = 3'b010, CFB = 3'b011, OFB = 3'b100
  input  logic [1:0]   aes_ctr_m_bits,       // AES CTR Increments: No inc = 2'b00, MSB 32 bits = 2'b01, MSB 64 bits = 2'b10, 128 bits = 2'b11
  input  logic	       DMA_Sync_Reset,
  input  logic [1:0]   i_key_size,
  output logic         aes_busy,             // 1= Aes core is busy
  output logic [127:0] aes_data_out,         // Data output bus to system
  output logic [127:0] aes_cipher_out,       // Direct block cipher output
  output logic [255:0] aes_working_key_debug,

  // Working Key Output
  output logic         lastround,
  output logic         load_expanded );

/**********************************************************************/
// Data Type Declarations (Signal Declarations)
/**********************************************************************/
  logic [127:0] work_in_IV;
  logic [127:0] work_out_IV;
  logic [127:0] eng_data_in;
  logic [127:0] eng_data_out;
  logic [127:0] eng_state_mask_in;
  logic [127:0] eng_new_mask_in;
  logic [127:0] eng_next_mask_out;
  logic	[127:0] working_mask;
  logic [127:0] statemask_in;
  logic [127:0] overlap_mask_in; // Consecutive rounds of encryption overlapping with lastround
  logic	[127:0]	prng_mask_out;
  logic [127:0] next_mask;

  logic 	encrypt;               // 1 => encrypt, 0 => decrypt
  logic 	mixcolumns;            // do the mix columns (or inverse for decrypt) operation
  logic 	bytesub_shiftrows;     // do the byte sub and shift rows (or inverse for decrypt) operations
  logic 	add_roundkey;          // do xor operation with round key 
  logic 	sbox_input_sel;        // input selector for the key SBOX
  logic 	store_key;             // store the key
  logic 	store_state;           // store the state
  logic [7:0] 	rcon;                  // the round constant used for key expansion
  logic [127:0] round_key;             // working key output from key scheduler 

  logic 	even_rounds;
  logic 	odd_rounds;

  logic       aes_start; 
  logic [127:0] work_in_iv_data;
  logic 	work_in_iv_en;
  logic [127:0] work_out_iv_data;
  logic 	work_out_iv_en;
  logic 	n_busy;
  logic 	aes_busy_s;
  logic 	aes_key_start;
  logic [127:0] aes_ctr_incr;

  // In CFB and OFB cipher always encrypts but data connections are different for decrypt
  // eng_aes_op drives the cipher and it is set always to encrypt in CFB and OFB modes
  logic         eng_aes_op; 


/**********************************************************************/
//Structural instantiations 
/**********************************************************************/

//B_FSM - resposible what's going in/out from engine.
  
  assign	data_valid = aes_key_start;

  assign        aes_cipher_out = eng_data_out;

  // AES Busy register
  always_ff @(posedge clk or negedge reset_n)
    begin
      if (~reset_n)
	  aes_busy_s <= 1'b0;
      else if(DMA_Sync_Reset)
          aes_busy_s <= 1'b0;
      else
	  aes_busy_s <= aes_busy;
    end

  assign	n_busy = aes_busy_s & ~aes_busy;


  // Mask Registers
  always_ff @(posedge clk or negedge reset_n)
    begin : internal_mask_registers
	  if (~reset_n)
	      working_mask	<= {128{1'b0}};

	  else if (aes_flush)
	      working_mask	<= {128{1'b0}};

	  else if (store_state | start)
	    begin
	      if (start)
		    working_mask	<= statemask_in;

	      else if (store_state)
		    working_mask	<= eng_next_mask_out;
  
	      else
		    working_mask 	<= working_mask;
	    end
    end

    assign statemask_in 	= (mask_logic_en) ? prng_mask_out : {128{1'b0}};

  // AES Counter bit increments
  always_comb
  begin
    case(aes_ctr_m_bits)
	       2'b00:begin
	       aes_ctr_incr        = work_in_IV;
	       end
	       2'b01:begin
	       aes_ctr_incr        = {work_in_IV[127:32], work_in_IV[31:0]+32'b1};
	       end
	       2'b10:begin
	       aes_ctr_incr        = {work_in_IV[127:64], work_in_IV[63:0]+64'b1};
	       end
	       2'b11:begin 
	       aes_ctr_incr        = {work_in_IV[127:0]+128'b1};
	       end
      endcase
    end


  // AES engine inputs/outputs
  always @*
    begin
      if (!reset_n)
      begin
        aes_data_out      = 'b0;
      end
      else
      begin
      case(aes_internal_mode)

		3'b000 : begin  
		  eng_data_in       = aes_data_in ^ statemask_in;
		  aes_data_out      = eng_data_out;
		  eng_state_mask_in = mask_logic_en ? working_mask : {128{1'b0}};
		  eng_new_mask_in   = mask_logic_en ? prng_mask_out : {128{1'b0}};
		  eng_next_mask_out = next_mask;
		  work_in_iv_data   = 128'b0;
		  work_in_iv_en     = 1'b0;
		  work_out_iv_data  = 128'b0;
		  work_out_iv_en    = 1'b0;
		end

		3'b001 : begin  
		  eng_data_in     = ~aes_op? (lastround & start) ? (eng_data_out ^ aes_data_in ^ statemask_in) : (work_in_IV ^ aes_data_in ^ statemask_in) : (aes_data_in ^ statemask_in);
		  aes_data_out     = aes_op? (eng_data_out ^ work_out_IV) : eng_data_out;
		  eng_state_mask_in = mask_logic_en ? working_mask : {128{1'b0}};
		  eng_new_mask_in   = mask_logic_en ? prng_mask_out : {128{1'b0}};
		  eng_next_mask_out = next_mask;
		  work_in_iv_data = ~aes_op? (aes_load_iv ? aes_iv_in : lastround ? eng_data_out : 128'b0) : aes_key_start ? aes_data_in : 128'b0;
		  work_in_iv_en   = aes_load_iv | lastround |(aes_op & aes_key_start);
		  work_out_iv_data =  aes_load_iv ? aes_iv_in : aes_op & lastround ? work_in_IV : 128'b0;
		  work_out_iv_en   = (lastround & aes_op) | aes_load_iv;
		end          

		3'b010 : begin  
		  eng_data_in       = work_in_IV ^ statemask_in;
		  aes_data_out      = eng_data_out ^ work_out_IV;
		  eng_state_mask_in = mask_logic_en ? working_mask  : {128{1'b0}};
		  eng_new_mask_in   = mask_logic_en ? prng_mask_out : {128{1'b0}};
		  eng_next_mask_out = next_mask;
		  work_in_iv_data   = aes_load_iv ? aes_iv_in : aes_key_start ? aes_ctr_incr : 128'b0;
		  work_in_iv_en     = aes_load_iv |aes_key_start;
		  work_out_iv_data  = aes_data_in;
		  work_out_iv_en    = aes_key_start;
		end

        3'b011 : begin  // CFB
		  eng_data_in     = work_in_IV ^ statemask_in;
		  aes_data_out     = eng_data_out ^ work_out_IV;
		  eng_state_mask_in = mask_logic_en ? working_mask : {128{1'b0}};
		  eng_new_mask_in   = mask_logic_en ? prng_mask_out : {128{1'b0}};
		  eng_next_mask_out = next_mask;
		  work_in_iv_data = aes_load_iv ? aes_iv_in : lastround? (~aes_op? eng_data_out ^ work_out_IV : work_out_IV) : 128'b0;
		  work_in_iv_en   = aes_load_iv | lastround;
		  work_out_iv_data = aes_data_in;
		  work_out_iv_en   = aes_key_start;
		end             

		3'b100 : begin  // OFB
		  eng_data_in     = work_in_IV ^ statemask_in;
		  aes_data_out     = eng_data_out ^ work_out_IV;
		  eng_state_mask_in = mask_logic_en ? working_mask : {128{1'b0}};
		  eng_new_mask_in   = mask_logic_en ? prng_mask_out : {128{1'b0}};
		  eng_next_mask_out = next_mask;
		  work_in_iv_data = aes_load_iv ? aes_iv_in : lastround? eng_data_out : 128'b0;
		  work_in_iv_en   = aes_load_iv | lastround;
		  work_out_iv_data = aes_data_in;
		  work_out_iv_en   = aes_key_start;
		end             

		3'b111:begin  
		  eng_data_in      = work_in_IV ^ aes_data_in ^ statemask_in;
		  aes_data_out     = eng_data_out ^ work_out_IV;
		  eng_state_mask_in = mask_logic_en ? working_mask : {128{1'b0}};
		  eng_new_mask_in   = mask_logic_en ? prng_mask_out : {128{1'b0}};
		  eng_next_mask_out = next_mask;
		  work_in_iv_data  = aes_iv_in;
		  work_in_iv_en    = aes_load_iv;
		  work_out_iv_data = aes_iv_in;
		  work_out_iv_en   = aes_load_iv;
		end
                default : begin
		  eng_data_in      = work_in_IV ^ aes_data_in ^ statemask_in;
		  aes_data_out     = eng_data_out ^ work_out_IV;
		  eng_state_mask_in = mask_logic_en ? working_mask : {128{1'b0}};
		  eng_new_mask_in   = mask_logic_en ? prng_mask_out : {128{1'b0}};
		  eng_next_mask_out = next_mask;
		  work_in_iv_data  = aes_iv_in;
		  work_in_iv_en    = aes_load_iv;
		  work_out_iv_data = aes_iv_in;
		  work_out_iv_en   = aes_load_iv;
               end
	endcase // case (aes_mode)
        end
      end // always_comb

  // Registers for IV
   always_ff @(posedge clk or negedge reset_n)
    if(~reset_n)
     work_in_IV <= 128'b0;
    else if(work_in_iv_en)  
     work_in_IV <= work_in_iv_data ;
 
  always_ff @(posedge clk or negedge reset_n)
    if(~reset_n)
     work_out_IV <= 128'b0;
    else if(work_out_iv_en)  
     work_out_IV <= work_out_iv_data ;
// Output of the working key

  // Module instantiations
  // PRNG Module

  A_cse_ocs_dpa_aes10c_xorshift i_xorshft (
	// INPUTS
	.seed		(aes_mask_seed),
	.reseed		(aes_load_mask),
	.clk		(clk),
	.reset_n	(reset_n),
	//Outputs
	.random_out	(prng_mask_out)
  );

  assign eng_aes_op = (aes_internal_mode == 3'b011) || (aes_internal_mode == 3'b100) ? 1'b0 : aes_op;


  A_cse_ocs_dpa_aes10c_control #(.AES_P(AES_P)) i_aes_control (
	//inputs
	.sys_clk           (clk),
	.reset_n           (reset_n),
	.data_valid 	   (start), 
	.aes_op            ({1'b0,eng_aes_op}),
	.i_key_size        (i_key_size),
	//outputs
	.aes_busy          (aes_busy),             
	.encrypt           (encrypt),              
	.mixcolumns        (mixcolumns),           
	.bytesub_shiftrows (bytesub_shiftrows),    
	.add_roundkey      (add_roundkey),
        .DMA_Sync_Reset	   (DMA_Sync_Reset),
	.store_state       (store_state),
	.store_key         (store_key),
	.sbox_input_sel    (sbox_input_sel),       
	.rcon              (rcon),
	.even_rounds       (even_rounds),			 
	.odd_rounds        (odd_rounds),
	.lastround         (lastround),
	.aes_start         (aes_start),
	.aes_key_start 	   (aes_key_start),
	.load_expanded     (load_expanded)		       
  );

  A_cse_ocs_dpa_aes10c_state #(.OCS_AES_DPA_USE_CTECH(OCS_AES_DPA_USE_CTECH),
                             .AES_P(AES_P))
     i_aes_state (
	//INPUTS
	  .sys_clk           (clk),
	  .reset_n           (reset_n),
	  .flush             (aes_flush),
	  .encrypt           (encrypt),
	  .mixcolumns        (mixcolumns),
	  .bytesub_shiftrows (bytesub_shiftrows),
	  .add_roundkey      (add_roundkey),      
	  .round_key         (round_key),
	  .store_state       (store_state),
	  .data_in           (eng_data_in),
	  .state_mask	     (eng_state_mask_in),
	  .new_mask	     (eng_new_mask_in),
	  .aes_start         (aes_key_start),
	  .lastround	     (lastround),
	  //outputs
	  .data_out          (eng_data_out),
	  .next_mask	     (next_mask)
  );

  A_cse_ocs_dpa_aes10c_key i_aes_key (
//inputs
	  .sys_clk           (clk),
	  .reset_n           (reset_n),
	  .flush             (aes_flush),
	  .key_in            (aes_key_in),
	  .store_key         (store_key),
	  .encrypt           (encrypt),
	  .sbox_input_sel    (sbox_input_sel),
	  .rcon              (rcon),
	  .aes_key_start     (aes_key_start),
	  .even_rounds       (even_rounds),
	  .odd_rounds        (odd_rounds),
	  .i_key_size        (i_key_size),			  
	  //outputs
	  .round_key         (round_key),
	  .aes_expand_key    (aes_working_key_debug)
);

endmodule
